from heapq import nlargest
from itertools import count, product
from time import time

import numpy as np
import utils_textrank as util
import zhtts
from faster_whisper import WhisperModel
from gensim.models import Word2Vec
from moviepy.editor import AudioFileClip, CompositeAudioClip, VideoFileClip


class FastTextRank4Sentence:
    def __init__(self, use_stopword=False, stop_words_file=None, use_w2v=False, dict_path=None, max_iter=100, tol=0.0001):
        """
        初始化FastTextRank4Sentence对象。

        参数:
        use_stopword (bool): 是否使用停用词过滤。
        stop_words_file (str): 停用词文件的路径。
        use_w2v (bool): 是否使用Word2Vec词向量模型。
        dict_path (str): 词向量模型的路径。
        max_iter (int): TextRank算法的最大迭代次数。
        tol (float): TextRank算法的收敛容差。
        """
        self.use_stopword = use_stopword
        self.use_w2v = use_w2v
        self.max_iter = max_iter
        self.tol = tol
        self.word2vec = Word2Vec.load(dict_path) if use_w2v and dict_path else None
        self.stop_words = self._load_stop_words(stop_words_file) if use_stopword else set()
        np.seterr(all='warn')

    @staticmethod
    def _load_stop_words(file_path):
        """
        加载停用词文件。

        参数:
        file_path (str): 停用词文件的路径。

        返回:
        set: 包含停用词的集合。
        """
        with open(file_path, 'r', encoding='utf-8') as f:
            return set(line.strip() for line in f)

    def filter_dictword(self, sents):
        """
        过滤掉不在词向量字典中的词汇。

        参数:
        sents (list): 分词后的句子列表。

        返回:
        list: 过滤后的句子列表。
        """
        return [
            [word for word in sentence if word in self.word2vec.wv]
            for sentence in sents
        ]

    def summarize(self, text, n=5):
        """
        对给定文本进行摘要提取。

        参数:
        text (str): 待摘要的文本。
        n (int): 需要提取的句子数量。

        返回:
        list: 提取出的摘要句子列表。
        """
        text = util.as_text(text.replace('\n', '').replace('\r', ''))
        sentences, sents = util.cut_filter_words(util.cut_sentences(text), self.stop_words, self.use_stopword)
        if self.use_w2v:
            sents = self.filter_dictword(sents)
        graph = self._create_graph_sentence(sents)
        scores = util.weight_map_rank(graph, self.max_iter, self.tol)
        selected_indices = [i for _, i in nlargest(n, zip(scores, count()))]
        return [sentences[i] for i in selected_indices]

    def _create_graph_sentence(self, word_sent):
        """
        创建句子相似度图（邻接矩阵）。

        参数:
        word_sent (list): 分词后的句子列表。

        返回:
        np.array: 句子之间相似度的邻接矩阵。
        """
        num = len(word_sent)
        board = np.zeros((num, num))
        for i, j in product(range(num), repeat=2):
            if i != j:
                board[i][j] = (
                    self._compute_similarity_by_avg(word_sent[i], word_sent[j])
                    if self.use_w2v else util.two_sentences_similarity(word_sent[i], word_sent[j])
                )
        return board

    def _compute_similarity_by_avg(self, sents_1, sents_2):
        """
        计算两个句子的平均词向量余弦相似度。

        参数:
        sents_1 (list): 第一个句子的分词列表。
        sents_2 (list): 第二个句子的分词列表。

        返回:
        float: 两个句子之间的相似度得分。
        """
        if not sents_1 or not sents_2:
            return 0.0
        vec1 = np.mean([self.word2vec.wv[word] for word in sents_1], axis=0)
        vec2 = np.mean([self.word2vec.wv[word] for word in sents_2], axis=0)
        return util.cosine_similarity(vec1, vec2)


def initialize_model(model_size: str, device: str, compute_type: str) -> WhisperModel:
    """
    初始化WhisperModel模型。

    参数:
    model_size (str): 模型的大小。
    device (str): 使用的设备（如"cuda"或"cpu"）。
    compute_type (str): 计算类型（如"float16"或"int8_float16"）。

    返回:
    WhisperModel: 初始化的WhisperModel对象。
    """
    return WhisperModel(model_size, device=device, compute_type=compute_type)


def transcribe_audio(model: WhisperModel, file_path: str, beam_size: int):
    """
    使用指定的模型对音频文件进行转录。

    参数:
    model (WhisperModel): 用于转录的WhisperModel对象。
    file_path (str): 音频或视频文件的路径。
    beam_size (int): 转录时使用的beam大小。

    返回:
    Tuple[list, dict]: 转录的段落和附加信息。
    """
    return model.transcribe(file_path, beam_size=beam_size)


def summarize_text(text: str, num_sentences: int = 3) -> str:
    """
    对文本进行摘要提取。

    参数:
    text (str): 待提取摘要的文本。
    num_sentences (int): 提取的句子数量。

    返回:
    str: 提取的摘要文本。
    """
    summarizer = FastTextRank4Sentence(use_stopword=True, stop_words_file='stopwords.txt', use_w2v=False)
    return ''.join(summarizer.summarize(text, n=num_sentences))



def text_to_speech(text: str, output_path: str):
    """
    将文本转换为语音。

    参数:
    text (str): 待转换的文本。
    output_path (str): 输出音频文件的路径。
    """
    tts = zhtts.TTS() # use fastspeech2 by default
    tts.text2wav(text=text, wavpath=output_path)


def sum_and_tts(segments, info, audio_path):
    """
    打印转录结果，包括检测到的语言、完整转录文本和摘要。

    参数:
    segments (list): 转录的段落列表。
    info (dict): 附加转录信息，包括语言检测结果。
    """
    print(f"Detected language '{info.language}' with probability {info.language_probability:.2f}")

    full_text = "。".join(segment.text for segment in segments)
    summary = summarize_text(full_text, 3)

    print("Full Transcription:")
    print(full_text)
    print("\nSummary:")
    print(summary)
    print("Text to speech:")
    text_to_speech(summary, audio_path)


def combine_audio_video(video_path: str, audio_path: str, output_path: str):
    """
    使用ffmpeg-python库将音频文件与视频文件合并为一个文件。

    参数:
    video_path (str): 视频文件的路径。
    audio_path (str): 音频文件的路径。
    output_path (str): 输出文件的路径。
    """
    video_input = VideoFileClip(video_path)
    audio_input = AudioFileClip(audio_path)
    video_with_audio = video_input.set_audio(audio_input)
    video_with_audio.write_videofile(output_path)
    print(f"Combined video saved to {output_path}")

def main():
    """
    主函数，执行模型初始化、音频转录、结果打印和音视频合并操作。
    """
    model_size = "large-v3"
    sample_name = 'douyin2'
    orig_video_path = f"./videos/{sample_name}.mp4"
    summary_video_path = f"./summary_videos/{sample_name}.mp4"
    output_audio_path = f"./final_output/{sample_name}.wav"
    output_video_path = f"./final_output/{sample_name}.mp4"

    start_time = time()

    model = initialize_model(model_size=model_size, device="cuda", compute_type="float16")
    print(f"Model initialization took {time() - start_time:.2f} seconds.")

    segments, info = transcribe_audio(model, orig_video_path, beam_size=5)
    print(f"Audio transcription took {time() - start_time:.2f} seconds.")

    sum_and_tts(segments, info, output_audio_path)
    print(f"Summarization and TTS took {time() - start_time:.2f} seconds.")

    # Combine audio and video
    combine_audio_video(summary_video_path, output_audio_path, output_video_path)

    print(f"Total runtime: {time() - start_time:.2f} seconds.")

if __name__ == "__main__":
    main()
